# coding : UTF-8
"""
作者：BingBO   时间：2022年10月23日
自动调整代码格式 ：Alt+Ctrl+L
"""
import cv2 as cv
import numpy as np
import math
import random
from numpy.linalg import solve  # 解线性方程


# from skimage import morphology  # 形态学操作，如开闭运算、骨架提取等


def cog_Roi(Droi, startP, a, mode):
    h, w = Droi.shape
    Py = startP[0]
    Px = startP[1]
    gravityCoorList = []
    if mode == "laserL1":
        for row in range(h - 1, -1, -1):
            n = 0
            sum_ColIdx = 0
            for col in range(w):
                if Droi[row, col] != 0:
                    n += 1
                    sum_ColIdx += col + 1
            if n:
                gravityCol = round(sum_ColIdx / n) - 1
                gravityRow = row
                # roi在起点上方
                if Px >= 2 * a:  # 需要将Roi坐标 转成 在原图中的坐标（行，列）
                    gravityCoorList.append((Py - a + gravityRow, Px - 2 * a + gravityCol))
                else:
                    gravityCoorList.append((Py - a + gravityRow, gravityCol))
    if mode == "laserL2":
        for row in range(h):
            n = 0
            sum_ColIdx = 0
            for col in range(w):
                if Droi[row, col] != 0:
                    n += 1
                    sum_ColIdx += col + 1
            if n:
                gravityCol = round(sum_ColIdx / n) - 1
                gravityRow = row
                # roi在起点下方
                if Px >= 2 * a:  # 需要将Roi坐标 转成 在原图中的坐标（行，列）
                    gravityCoorList.append((Py + gravityRow + 1, Px - 2 * a + gravityCol))
                else:
                    gravityCoorList.append((Py + gravityRow + 1, gravityCol))
    return gravityCoorList


# 纵向 灰度重心法
def cog(binary, mode, show=None):
    height, width = binary.shape
    print("cog_binary.shape:", binary.shape)
    cp_row, cp_col = [], []
    if mode == 0:  # 垂直COG
        for col in range(width):
            x0, y0 = 0, 0
            for row in range(height):
                x0 += binary[row, col]
                y0 += binary[row, col] * (row + 1)
            if x0 != 0:
                indexRow = round(y0 / x0) - 1
                cp_row += [indexRow]
                cp_col += [col]  # 中心点列索引(包含空白区域列坐标的突变)
    elif mode == 1:  # 水平COG
        for row in range(height):
            x0, y0 = 0, 0
            for col in range(width):
                x0 += binary[row, col]
                y0 += binary[row, col] * (col + 1)
            if x0 != 0:
                indexRow = round(y0 / x0) - 1
                cp_col += [indexRow]
                cp_row += [row]  # 中心点列索引(包含空白区域列坐标的突变)
    # print("CP_row:", cp_row)
    # print("CP_col:", cp_col)
    if show == 1:
        # cpGlobal = np.zeros_like(binary)
        cpShow = np.zeros([height, width, 3], np.uint8)
        cpShow[:, :, 0] = binary
        cpShow[:, :, 1] = binary
        cpShow[:, :, 2] = binary
        for i in range(0, len(cp_row)):
            cpShow[cp_row[i], cp_col[i]] = [0, 0, 255]  # 红色
        cv.imshow("cpShow", cpShow)
        cv.imwrite(r"F:\MVS_Data\cpShow.png", cpShow)
    return cp_row, cp_col


def vPointlocation(binaryWeldment, mode):
    # 传入的是  边缘点左侧的 光条阈值图像
    if mode == 1:
        mode_str = "laser1_"
    elif mode == 2:
        mode_str = "laser2_"

    cp_row, cp_col = cog(binaryWeldment, 1, show=1)  # 水平COG （为了确保垂直坡口处的中心点数量）
    # cp_row, cp_col = roiIteration(binaryWeldment, CrossPoint, 40)  # 迭代ROI获取中心点

    # 创建中心点图像
    cpImg = np.zeros_like(binaryWeldment)
    for i in range(0, len(cp_row)):
        cpImg[cp_row[i], cp_col[i]] = 255
    # cv.imwrite("F:\\MVS_Data\\cpImg_" + mode_str + ".png", cpImg)

    # 判断光条类型
    numAll = 0
    if mode == 1:
        for i in range(8, len(cp_row) - 8):
            row = cp_row[i]
            col = cp_col[i]
            num = 0
            # 统计右下
            for m in range(0, 4):  # 0123
                for n in range(1, 8):  # 1~7
                    if cpImg[row + n, col + m]:  # 会超出图像的大小
                        num += 1
            Idx = [[1, 1], [1, 2], [1, 3], [2, 1], [2, 2], [2, 3], [3, 2], [3, 3], [4, 3]]
            for k in range(len(Idx)):
                if cpImg[row + Idx[k][0], col + Idx[k][1]]:
                    num -= 1  # 减去这些黑点
            # 统计左上
            Idx = [[0, -1], [-1, -1], [-1, -2], [-1, -3], [-2, -2], [-2, -3], [-2, -4], [-2, -5], [-3, -3], [-3, -4],
                   [-3, -5], [-3, -6]]
            for k in range(len(Idx)):
                if cpImg[row + Idx[k][0], col + Idx[k][1]]:
                    num += 2
            if num > numAll:
                numAll = num
                Index = i - 1
                vPoint = (cp_row[Index], cp_col[Index])  # 实验证明，选取上一点的精度更高
    if mode == 2:
        for i in range(8, len(cp_row) - 8):
            row = cp_row[i]
            col = cp_col[i]
            num = 0
            # 统计右上
            for m in range(0, 4):  # 0123
                for n in range(1, 8):  # 1~7
                    if cpImg[row - n, col + m]:
                        num += 1
            Idx = [[-1, 1], [-1, 2], [-1, 3], [-2, 1], [-2, 2], [-2, 3], [-3, 2], [-3, 3], [-4, 3]]
            for k in range(len(Idx)):
                if cpImg[row + Idx[k][0], col + Idx[k][1]]:
                    num -= 1
            # 统计左下
            Idx = [[0, -1], [1, -1], [1, -2], [1, -3], [2, -2], [2, -3], [2, -4], [2, -5], [3, -3], [3, -4], [3, -5],
                   [3, -6]]
            for k in range(len(Idx)):
                if cpImg[row + Idx[k][0], col + Idx[k][1]]:
                    num += 2
            if num > numAll:
                numAll = num
                Index = i - 1
                vPoint = (cp_row[Index], cp_col[Index])  # 实验证明，选取上一点的精度更高
    print(mode_str + " V点定位坐标：", vPoint)
    # print("该点的统计点数：", numAll)

    # 显示定位精度 图像
    h, w = binaryWeldment.shape[:2]
    cpImg3 = np.zeros([h, w, 3], np.uint8)
    cpImg3[:, :, 0] = cpImg
    cpImg3[:, :, 1] = cpImg
    cpImg3[:, :, 2] = cpImg
    cpImg3[vPoint[0], vPoint[1]] = [0, 0, 255]
    cv.imwrite("F:\\MVS_Data\\V-point—Accuracy_" + mode_str + ".png", cpImg3)

    # ***********分割最远点左右点集*************
    cp_rowL, cp_colL = cp_row[:Index + 1], cp_col[:Index + 1]
    cp_rowR, cp_colR = cp_row[Index + 1:], cp_col[Index + 1:]

    cpL = [cp_rowL, cp_colL]
    cpR = [cp_rowR, cp_colR]
    return cpL, cpR


def RANSAC(pointsRow, pointsCol):
    """
    .   @param pointsRow, pointsCol：点集
    .   @retval Coef_Mainline：拟合直线的Ax+By+C=0的系数 ABC
    """
    print("---------RANSAC ()主线系数 ---------")
    SIZE = len(pointsRow)
    print("总点数:", SIZE)
    ITERATION = SIZE * 2
    threshold = 0.9
    best_k, best_b = 0, 0
    innerVertical, innerSlope = 0, 0
    num = 0
    for i in range(ITERATION):
        num += 1
        sampleIndex = random.sample(range(SIZE), 2)  # 每次迭代都 随机截取两个样本，作为数据索引值
        # 随机取得的两个坐标点(x1，y1)、(x2，y2)
        x1, y1 = pointsCol[sampleIndex[0]], pointsRow[sampleIndex[0]]
        x2, y2 = pointsCol[sampleIndex[1]], pointsRow[sampleIndex[1]]
        if x2 != x1:
            k = (y2 - y1) / (x2 - x1)  # 此处可以改进一下
            b = y1 - k * x1
            innerP1 = 0
            for index in range(SIZE):  # 遍历所有数据点
                x, y = pointsCol[index], pointsRow[index]
                d = math.fabs(k * x - y + b) / math.sqrt(1 + k ** 2)
                if d <= 1:  # 1/sqrt2
                    innerP1 += 1
            if innerP1 > innerSlope:
                innerSlope = innerP1
                best_k, best_b = k, b
            if innerSlope > SIZE * threshold:
                break
        else:  # x2 == x1,存在随机不到两个数相同的概率
            innerP2 = 0
            for index in range(SIZE):
                x, y = pointsCol[index], pointsRow[index]
                if x == x1:
                    innerP2 += 1
            if innerP2 > innerVertical:
                innerVertical = innerP2
                C = -x2
    print("迭代 %d 次" % num)
    if innerVertical >= innerSlope:
        A, B, C = 1, 0, C
        print("垂直_最大内点数:{}, 占总点数:{:.2f}%".format(innerVertical, innerVertical / SIZE * 100))
    else:
        A, B, C = best_k, -1, best_b
        print("倾斜_最大内点数:{}, 占总点数:{:.2f}%".format(innerSlope, innerSlope / SIZE * 100))
    LineCoef = [A, B, C]
    print("直线系数 A = {:.4f}, B = {}, C = {:.4f}".format(A, B, C))
    return LineCoef


def pointSegment(pointSet):
    print("\n---------正在调用 pointSegment()---------")
    pointSet_row, pointSet_col = pointSet
    # RANSAC
    LineCoef = RANSAC(pointSet_row, pointSet_col)
    best_k, B, best_b = LineCoef  # A1, B1, C1 = best_k, -1, best_b

    crossLineRow, crossLineCol = [], []
    grooveRow, grooveCol = [], []
    countLine, countGroove = 0, 0

    if 0.9 < math.fabs(best_k) <= 1 and B == -1:  # 排除垂直A=1的情况
        print("工件点集系数：", LineCoef)
        for i in range(0, len(pointSet_col)):
            d = math.fabs(best_k * pointSet_col[i] - pointSet_row[i] + best_b) / math.sqrt(1 + best_k ** 2)
            if d <= 1:
                crossLineRow += [pointSet_row[i]]
                crossLineCol += [pointSet_col[i]]
                countLine += 1
            if d > 3:
                grooveRow += [pointSet_row[i]]
                grooveCol += [pointSet_col[i]]
                countGroove += 1
    elif best_k == 1 and B == 0:
        # 垂直：math.fabs(best_k) = 1, B=0
        print("坡口(垂直)点集系数：", LineCoef)
        for i in range(0, len(pointSet_col)):
            d = math.fabs(best_k * pointSet_col[i] - pointSet_row[i] + best_b) / math.sqrt(1 + best_k ** 2)
            if d == 0:
                grooveRow += [pointSet_row[i]]
                grooveCol += [pointSet_col[i]]
                countGroove += 1
            if d > 3:
                crossLineRow += [pointSet_row[i]]
                crossLineCol += [pointSet_col[i]]
                countLine += 1
    else:
        # 大斜率：1 < math.fabs(best_k), B=-1
        # 小斜率：math.fabs(best_k) < 0.9, B=-1
        print("坡口(倾斜)点集系数：", LineCoef)
        for i in range(0, len(pointSet_col)):
            d = math.fabs(best_k * pointSet_col[i] - pointSet_row[i] + best_b) / math.sqrt(1 + best_k ** 2)
            if d <= 1:
                grooveRow += [pointSet_row[i]]
                grooveCol += [pointSet_col[i]]
                countGroove += 1
            if d > 3:
                crossLineRow += [pointSet_row[i]]
                crossLineCol += [pointSet_col[i]]
                countLine += 1
    print("坡口%d个点" % countGroove)
    print("非坡口光条%d个点" % countLine)
    groovePoints = [grooveRow, grooveCol]
    crossLinePoints = [crossLineRow, crossLineCol]
    return groovePoints, crossLinePoints


# 最小二乘法拟合直线
def fitLine(pointSet):
    pointSet_row, pointSet_col = pointSet
    # 垂直的情况下
    if max(pointSet_col) == min(pointSet_col):
        A, B, C = 1, 0, -pointSet_col[0]
    else:
        poly = np.polyfit(pointSet_col, pointSet_row, 1)  # 用最小二乘法（含COG_V误差）拟合得到的直线参数
        fit_k, fit_b = poly[0], poly[1]
        A, B, C = fit_k, -1, fit_b
    LineCoef = [A, B, C]
    return LineCoef


def LaserExtract(src, mode):
    # src ：传入的已经是 边缘点左边的 工件光条原始图像
    Red = src[:, :, 2]  # B0 G1 R2
    _, threshold = cv.threshold(Red, 190, 255, cv.THRESH_BINARY)  # 转为二值图
    # threshold2 = cv.adaptiveThreshold(Red, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 13, 2)
    threshold = cv.medianBlur(threshold, 3)

    if mode == 1:
        mode_str = "Laser1_"
        print("=========================== laser1 特征点提取 ===========================")
    elif mode == 2:
        mode_str = "Laser2_"
        print("=========================== laser2 特征点提取 ===========================")
    # 左右中心点集
    cpL, cpR = vPointlocation(threshold, mode)  # 不再是迭代ROI 而是全局阈值 再cog提取中心点
    # V点左侧 点击分割
    groovePointsL, crossLinePointsL = pointSegment(cpL)  # 调用RANSAC 实现分割
    # V点右侧 点击分割
    groovePointsR, crossLinePointsR = pointSegment(cpR)  # 调用RANSAC 实现分割

    Coef_GrooveL = fitLine(groovePointsL)
    Coef_WeldL = fitLine(crossLinePointsL)
    Coef_GrooveR = fitLine(groovePointsR)
    Coef_WeldR = fitLine(crossLinePointsR)

    Al, Bl, Cl = Coef_WeldL
    Agl, Bgl, Cgl = Coef_GrooveL
    Agr, Bgr, Cgr = Coef_GrooveR
    Ar, Br, Cr = Coef_WeldR

    print("------" + mode_str + "3个点计算结果" + "------")
    # 计算3个特征点坐标
    A1 = np.array([[Al, Bl], [Agl, Bgl]])  # 输出系数矩阵A
    b1 = np.array([-Cl, -Cgl])  # 值
    fp_x1, fp_y1 = solve(A1, b1)
    fp_y1, fp_x1 = round(fp_y1), round(fp_x1)
    # cpImg[fp_y1, fp_x1] = [255, 0, 255]
    cv.circle(src, [fp_x1, fp_y1], 1, (0, 255, 0), -1)
    cv.circle(src, [fp_x1, fp_y1], 5, (255, 0, 255), 1)
    P1 = [fp_x1, fp_y1]
    print(mode_str + "P1：", P1)

    A2 = np.array([[Agl, Bgl], [Agr, Bgr]])  # 输出系数矩阵A
    b2 = np.array([-Cgl, -Cgr])  # 值
    fp_x2, fp_y2 = solve(A2, b2)
    fp_y2, fp_x2 = round(fp_y2), round(fp_x2)
    cv.circle(src, [fp_x2, fp_y2], 1, (0, 255, 0), -1)
    cv.circle(src, [fp_x2, fp_y2], 5, (255, 0, 255), 1)
    Pmid = [fp_x2, fp_y2]
    print(mode_str + "Pmid：", Pmid)

    A3 = np.array([[Agr, Bgr], [Ar, Br]])  # 输出系数矩阵A
    b3 = np.array([-Cgr, -Cr])
    fp_x3, fp_y3 = solve(A3, b3)
    fp_y3, fp_x3 = round(fp_y3), round(fp_x3)
    cv.circle(src, [fp_x3, fp_y3], 1, (0, 255, 0), -1)
    cv.circle(src, [fp_x3, fp_y3], 5, (255, 0, 255), 1)
    P2 = [fp_x3, fp_y3]
    print(mode_str + "P2：", P2)
    cv.imshow('FP_' + mode_str, src)
    cv.imwrite("F:\\MVS_Data\\FP_" + mode_str + ".png", src)

    return Pmid, P1, P2


def roundUp(value):
    # 返回value的整数部分与小数部分，两部分的数值符号与x相同，整数部分以浮点型表示
    a, b = math.modf(value)  # a:小数   b：整数
    if a >= 0.5:
        return int(b + 1)
    if a < 0.5:
        return int(b)


# 在图像上绘制特征线
def featureLineDisplay(img, P1mid, P2mid):
    featureP11, featureP12, featureP13 = P1mid  # laserL1  右下
    featureP21, featureP22, featureP23 = P2mid  # laserL2  右上

    x13, x14, x15 = featureP11[0], featureP12[0], featureP13[0]
    y13, y14, y15 = featureP11[-1], featureP12[-1], featureP13[-1]
    x23, x24, x25 = featureP21[0], featureP22[0], featureP23[0]
    y23, y24, y25 = featureP21[-1], featureP22[-1], featureP23[-1]

    cv.line(img, (x13, y13), (x14, y14), (34, 139, 34), 1)  # 绿色
    cv.line(img, (x14, y14), (x15, y15), (255, 144, 30), 1)  # 粉红

    cv.line(img, (x23, y23), (x24, y24), (34, 139, 34), 1)  # 绿色
    cv.line(img, (x24, y24), (x25, y25), (255, 144, 30), 1)  # 粉红

    return img
